package com.sun.handler;

import com.sun.config.GlobalVariableConfig;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import javax.annotation.PostConstruct;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

import static com.sun.constants.Constant.*;

@Slf4j
@Component
public class MyWebSocketHandler extends TextWebSocketHandler {

    private static final Map<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

    private static final Map<String, Integer> heartbeatCounters = new ConcurrentHashMap<>();

    private static final Map<String, String> pushMessageCache = new ConcurrentHashMap<>();

    @Value("${ws.server.heartbeat.clientSwitch:false}")
    private boolean clientHeartbeatSwitch;

    @Value("${ws.server.heartbeat.interval:10}")
    private int heartbeatInterval;

    @Value("${ws.server.heartbeat.maxMissCount:3}")
    private int heartbeatMaxMissCount;

    private final GlobalVariableConfig sessionCount;

    public MyWebSocketHandler(GlobalVariableConfig globalVariableConfig) {
        this.sessionCount = globalVariableConfig;
    }

    @PostConstruct
    public void initHeartbeat() {
        if(clientHeartbeatSwitch) {
            ScheduledExecutorService heartbeatExecutor = Executors.newScheduledThreadPool(1);
            heartbeatExecutor.scheduleAtFixedRate(() ->
                    broadcastMessage(HEART_BEAT, PING_MESSAGE), 0, heartbeatInterval, TimeUnit.SECONDS);
        }
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        heartbeatCounters.put(session.getId(), 0);
        sessions.put(session.getId(), session);
        sessionCount.incrementSessionCount();
        log.info("【WebSocket服务】连接成功，总连接数：{}", sessions.size());
        if(session.isOpen() && pushMessageCache.size() > 0) {
            for (Map.Entry<String, String> entry : pushMessageCache.entrySet()) {
                session.sendMessage(new TextMessage(entry.getValue()));
            }
        }
    }

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        if(PONG_MESSAGE.equalsIgnoreCase(message.getPayload())) {
            resetHeartbeatCounter(session);
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        closeSession(session);
        log.info("【WebSocket服务】连接断开，总连接数：{}", sessions.size());
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        closeSession(session);
        log.error("【WebSocket服务】发生异常：", exception);
    }

    private void resetHeartbeatCounter(WebSocketSession session) {
        heartbeatCounters.put(session.getId(), 0);
    }

    private void closeSession(WebSocketSession session) {
        if(session.isOpen()) {
            try {
                session.close();
            } catch (IOException e) {
                log.error("error", e);
            }
        }
        cleanupSession(session);
    }

    private void cleanupSession(WebSocketSession session) {
        heartbeatCounters.remove(session.getId());
        sessions.remove(session.getId());
        sessionCount.setSessionCount(sessions.size());
    }

    public synchronized void broadcastMessage(String type, String message) {
        sessions.forEach((sessionId, session) -> {
            try {
                if(!session.isOpen()) {
                    return;
                }
                session.sendMessage(new TextMessage(message));
                if(PING_MESSAGE.equalsIgnoreCase(message)) {
                    incrementHeartbeatCounter(session);
                    int heartbeatCounter = getHeartbeatCounter(session);
                    if(heartbeatCounter >= heartbeatMaxMissCount) {
                        log.info("【WebSocket服务】客户端心跳检查超过{}次，断开连接。", heartbeatCounter);
                        closeSession(session);
                    }
                } else {
                    pushMessageCache.put(type, message);
                }
            } catch (Exception e) {
                log.error("WebSocket server broadcast message error : ", e);
                closeSession(session);
            }
        });
    }

    private void incrementHeartbeatCounter(WebSocketSession session) {
        heartbeatCounters.computeIfPresent(session.getId(), (id, count) -> count + 1);
    }

    private int getHeartbeatCounter(WebSocketSession session) {
        return heartbeatCounters.getOrDefault(session.getId(), 0);
    }

}